#  Copyright 2022 The Open Islands Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import typing

from pyoi.constant import ComputingEngine, DataStructure, FederatedMode
from pyoi.context.abc import ContextABC, MQManagerABC
from pyoi.context.mq import MQ
from pyoi.context.types import MQArgs
from pyoi.runtime.types import RuntimeConfig
from pyoi.runtime_env import RuntimeENV
from pyoi.util import log_utils, IdUtils

LOGGER = log_utils.getLogger()


class Session:
    _SESSION: 'Session' = None

    @classmethod
    def get_or_create(cls, task_id: str, runtime_config: RuntimeConfig):
        mq_args, mq_manager = MQ.create(
            group=task_id,
            address=runtime_config.serviceAddresses.get("mq", {}),
            params=runtime_config.serviceParams.get("mq", {}),
        )
        if Session._SESSION is not None:
            return Session._SESSION
        Session._SESSION = Session(
            task_id=task_id,
            federated_mode=runtime_config.federatedMode,
            role=runtime_config.role,
            node_id=runtime_config.nodeId,
            nodes=runtime_config.nodes,
            mq_args=mq_args,
            mq_manager=mq_manager
        )
        return Session._SESSION

    @classmethod
    def get(cls) -> 'Session':
        return cls._SESSION

    def __init__(
            self,
            task_id: str,
            federated_mode: FederatedMode,
            role: str,
            node_id: str,
            nodes: dict,
            mq_args: MQArgs,
            mq_manager: MQManagerABC
    ):
        self._task_id = task_id
        self._session_id = IdUtils.get_task_session_id(task_id, role, node_id)
        self._federated_mode = federated_mode
        self._role = role
        self._node_id = node_id
        self._nodes = nodes
        self._mq_args = mq_args
        self._mq_manager = mq_manager

        self._contexts: typing.Dict[DataStructure, ContextABC] = {}
        self._default_context: ContextABC = None

    @property
    def task_id(self) -> str:
        return self._task_id

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def nodes(self):
        return self._nodes

    def _open(self):
        return self

    def _close(self):
        # todo: close all context
        pass

    def __enter__(self):
        return self._open()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_tb:
            LOGGER.exception("", exc_info=(exc_type, exc_val, exc_tb))
        return self._close()

    def create_context(self, data_structure: DataStructure):
        if data_structure in self._contexts:
            raise RuntimeError(f"{data_structure} already existed, please using get or create")
        if data_structure == DataStructure.SPARK_RDD:
            if RuntimeENV.COMPUTING_ENGINE == ComputingEngine.SPARK:
                from pyoi.context.rdd.spark.spark_rdd_context import SparkRDDContext
                context_id = "-".join([self._session_id, data_structure.name])
                context = SparkRDDContext(context_id=context_id, task_id=self._task_id,
                                          federated_mode=self._federated_mode,
                                          role=self._role, node_id=self._node_id, nodes=self._nodes,
                                          mq_args=self._mq_args,
                                          mq_manager=self._mq_manager, default_storage_engine=RuntimeENV.STORAGE_ENGINE)
                return context
            raise RuntimeError(f"{RuntimeENV.COMPUTING_ENGINE} not supported")
        raise RuntimeError(f"{data_structure} not supported")
